#include "Environment.h"
#define forEachDstr for (int i = 0; i < numDstr; i++)
#define forEachUstr for (int i = 0; i < numUstr; i++)
#define forEachBank for (int i = 0; i < numBanks; i++)
#define ONSCREEN 1
using namespace std;

// Constructor; initialize variables, allocates memory.
Environment::Environment (int n, double in[]) {
    t = 0;

    input = new double[n];
    for (int i = 0; i < n; i++) input[i] = in[i];
    rndGen = new RndGen;

    numDstr = static_cast <int> (input[0]);
    numUstr = static_cast <int> (input[1]);
    numBanks = static_cast <int> (input[2]);
    tMax = static_cast <int> (input[3]);

    ustrToLook = static_cast <int> (input[13]);
    banksToLook = static_cast <int> (input[14]);
    eps = input[15];

    // Instantianting the arrays of firms and banks
    dFirmsArray = new DownstreamFirm[numDstr];
    uFirmsArray = new UpstreamFirm[numUstr];
    banksArray = new Bank[numBanks];

    // Passing parameters to each one of the agents
    forEachDstr dFirmsArray[i].initialize (input);
    forEachUstr uFirmsArray[i].initialize (input);
    forEachBank banksArray[i].initialize (input);

    // Setup initial partners scheme
    forEachDstr {
        rndAgent = rndGen->getRndInt (0, numUstr - 1);
        dFirmsArray[i].setUstrPartner (&uFirmsArray[rndAgent]);
        uFirmsArray[rndAgent].addDstrPartner (&dFirmsArray[i]);

        rndAgent = rndGen->getRndInt (0, numBanks - 1);
        dFirmsArray[i].setBankPartner (&banksArray[rndAgent]);
        banksArray[rndAgent].addDstrPartner (&dFirmsArray[i]);
    }
    forEachUstr {
        rndAgent = rndGen->getRndInt (0, numBanks - 1);
        uFirmsArray[i].setBankPartner (&banksArray[rndAgent]);
        banksArray[rndAgent].addUstrPartner (&uFirmsArray[i]);
    }

    // Initialize output-related variables
    dstrAggrProdOut.open ("dstrAggrProd.dat");
    dstrSizeDistOut.open ("dstrSizeDist.dat");
    dstrAvgProfitOut.open ("dstrAvgProfit.dat");
    dstrFailOut.open ("dstrFail.dat");
    bankAvgIROut.open ("bankAvgIR.dat");
    bankIRDistOut.open ("bankIRDistOut.dat");
    bankAvgBDOut.open ("bankAvgBD.dat");
    bankFailOut.open ("bankFail.dat");
    bankEntropyOut.open ("bankEntropy.dat");
}

void // Running the simulation
Environment::start () {
    while (t < tMax) {
        step ();
        printOutput ();
        t++;
    }
}

void // One time step of the simulation
Environment::step () {
    if (ONSCREEN == 1) cout << "t = " << t << "\n";

    // Agents updates their indicators and quantities to be bought and sold
    forEachDstr {
        rndPrice = rndGen->getRndDouble (0, 1);
        dFirmsArray[i].calcQuantities (rndPrice);
    }
    forEachUstr uFirmsArray[i].calcQuantities ();
    forEachBank banksArray[i].calcQuantities ();

    // Firms calculate interest rates taxed on trades and loans
    forEachUstr uFirmsArray[i].calcInterestRates ();
    forEachDstr dFirmsArray[i].calcInterestRates ();

    // Agents calculate profit and update their net worths; if agent is insolvent
    // another one takes its place
    dstrFail = 0;
    forEachDstr {
        dFirmsArray[i].calcNetWorth ();
        if (dFirmsArray[i].checkInsolvency () == true) {
            dFirmsArray[i].goBankrupt ();
            dFirmsArray[i].initialize (input);
            dstrFail++;
        }
    }
    forEachUstr {
        uFirmsArray[i].calcNetWorth ();
        if (uFirmsArray[i].checkInsolvency () == true) {
            uFirmsArray[i].goBankrupt ();
            uFirmsArray[i].initialize (input);
        }
    }
    bankFail = 0;
    forEachBank {
        banksArray[i].calcNetWorth ();
        if (banksArray[i].checkInsolvency () == true) {
            banksArray[i].goBankrupt ();
            banksArray[i].initialize (input);
            bankFail++;
        }
    }

    // Firms updates their partners (preferred-partner choice)
    int newAgent;
    double minRate, rate;
    forEachDstr {
        newAgent = -1;
        minRate = dFirmsArray[i].getTradeIR ();
        probRandom = rndGen->getRndInt (1, 100) / 100;
        if (probRandom > (1 - eps) || dFirmsArray[i].gotNoUstrPartner ())
            newAgent = rndGen->getRndInt (0, numUstr - 1);
        else for (int j = 0; j < ustrToLook; j++) {
            rndAgent = rndGen->getRndInt (0, numUstr - 1);
            rate = uFirmsArray[rndAgent].getTradeIR ();
            if (rate < minRate) {
                minRate = rate;
                newAgent = rndAgent;
            }
        }

        if (newAgent != -1) {
            dFirmsArray[i].setUstrPartner (&uFirmsArray[newAgent]);
            uFirmsArray[newAgent].addDstrPartner (&dFirmsArray[i]);
        }

        newAgent = -1;
        minRate = dFirmsArray[i].getBankIR ();
        probRandom = rndGen->getRndInt (1, 100);
        if (probRandom > (1 - eps) || dFirmsArray[i].gotNoBankPartner ())
            newAgent = rndGen->getRndInt (0, numBanks - 1);
        else for (int j = 0; j < banksToLook; j++) {
            rndAgent = rndGen->getRndInt (0, numBanks - 1);
            rate = banksArray[rndAgent].getBankIR (dFirmsArray[i].getLeverage ());
            if (rate < minRate) {
                minRate = rate;
                newAgent = rndAgent;
            }
        }

        if (newAgent != -1) {
            dFirmsArray[i].setBankPartner (&banksArray[newAgent]);
            banksArray[newAgent].addDstrPartner (&dFirmsArray[i]);
        }
    }
    forEachUstr {
        newAgent = -1;
        minRate = dFirmsArray[i].getBankIR ();
        probRandom = rndGen->getRndInt (1, 100);
        if (probRandom > (1 - eps) || uFirmsArray[i].gotNoBankPartner ())
            newAgent = rndGen->getRndInt (0, numBanks - 1);
        else for (int j = 0; j < banksToLook; j++) {
            rndAgent = rndGen->getRndInt (0, numBanks - 1);
            rate = banksArray[rndAgent].getBankIR (uFirmsArray[i].getLeverage ());
            if (rate < minRate) {
                minRate = rate;
                newAgent = rndAgent;
            }
        }

        if (newAgent != -1) {
            uFirmsArray[i].setBankPartner (&banksArray[newAgent]);
            banksArray[newAgent].addUstrPartner (&uFirmsArray[i]);
        }
    }

}

void // Print desired output to data files
Environment::printOutput () {

    // Calculate and print downstream firms aggregate production
    dstrAggrProd = 0;
    forEachDstr dstrAggrProd += dFirmsArray[i].getProduction ();
    dstrAggrProdOut << t << "\t" << dstrAggrProd << "\n";

    // Print downstream firms size at t = 1000
    if (t == 1000) {
        double nw[numUstr];
        forEachUstr nw[i] = uFirmsArray[i].getNumPartners ();
        gsl_sort (nw, 1, numUstr);

        //for (int j = 0; j <= 1000; j++) nwHist[j] = 0;
        //forEachDstr
        //    for (int j = 0; j <= 1000; j++)
        //      if (nw[i] > j * nwMax/1000) nwHist[j]++;

        forEachUstr dstrSizeDistOut << nw[i] << "\t" << numUstr - i + 1 << "\n";
    }

    // Calculate and print downstream firms average profit
    dstrAvgProfit = 0;
    forEachDstr dstrAvgProfit += dFirmsArray[i].getProfit ();
    dstrAvgProfit /= numDstr;
    dstrAvgProfitOut << t << "\t" << dstrAvgProfit << "\n";

    // Print number of firm bankruptcies per turn
    dstrFailOut << t << "\t" << dstrFail << "\n";

    // Calculate and print average interest rate taxed by banks
    bankAvgIR = 0;
    forEachDstr bankAvgIR += dFirmsArray[i].getBankIR ();
    bankAvgIR /= numDstr;
    bankAvgIROut << t << "\t" << bankAvgIR << "\n";

    // Print banks IR distribution at t = 1000
    if (t == 1000) {
        double ir[numDstr], irMax, irHist[101];
        forEachDstr ir[i] = dFirmsArray[i].getBankIR ();
        irMax = gsl_stats_max (ir, 1, numDstr);

        for (int j = 0; j <= 100; j++) irHist[j] = 0;
        forEachDstr
            for (int j = 0; j <= 100; j++)
              if (ir[i] > j * irMax/100 && ir[i] < (j+1) * irMax/100) irHist[j]++;

        for (int j = 0; j <= 100; j++)
            bankIRDistOut << j * irMax/100 << "\t" << irHist[j] << "\n";
    }

    // Calculate and print average bank bad debt
    bankAvgBD = 0;
    forEachBank bankAvgBD += banksArray[i].getBadDebt ();
    bankAvgBD /= numBanks;
    bankAvgBDOut << t << "\t" << bankAvgBD << "\n";

    // Print number of bank bankruptcies per turn
    bankFailOut << t << "\t" << bankFail << "\n";

    // Calculate and print entropy of the bank system
    double totalNw, p;
    totalNw = 0;
    forEachDstr totalNw += dFirmsArray[i].getNetWorth ();
    bankEntropy = 0;
    forEachDstr {
        p = dFirmsArray[i].getNetWorth () / totalNw;
        bankEntropy -= p * log (p);
    }
    bankEntropy /= log (numDstr);
    bankEntropyOut << t << "\t" << 1 - bankEntropy << "\n";
}

//// Some other definitions and methods
// Auxiliar class RndGen constructor
RndGen::RndGen () {
    gsl_rng_env_setup ();
    T = gsl_rng_default;
    r = gsl_rng_alloc (T);

    // Change generator default seed
    int seed = 1001;
    gsl_rng_set (r, seed);
}

double // Returns a double between minsp and maxsp
RndGen::getRndDouble (double minsp, double maxsp) {
    return gsl_ran_flat (r, minsp, maxsp);
}

int // Returns a int between min and max
RndGen::getRndInt (int min, int max) {
    return min + gsl_rng_uniform_int (r, max - min + 1);
}
